
#include "SimpleMQTT.hpp"

namespace SimpleMQTT
{
    MQTTClient::~MQTTClient()
    {
        this->deinit();
    }
    void appendUTF8String(
        uint8_t *buffer,
        uint16_t *offset,
        const char *str)
    {
        uint16_t len = strlen(str);
        buffer[(*offset)] = len >> 8;
        buffer[(*offset)+1] = len & 0xFF;
        memcpy(&buffer[(*offset)+2],str,len);
        *offset += (len+2);
    }

    void appendByte(
        uint8_t *buffer,
        uint16_t *offset,
        uint8_t data)
    {
        buffer[(*offset)] = data;
        *offset += 1;
    }

    void appendHalfWord(
        uint8_t *buffer,
        uint16_t *offset,
        uint16_t data)
    {
        buffer[(*offset)] = data >> 8;
        buffer[(*offset)+1] = data & 0xFF;
        *offset += 2;
    }

    void appendFixedMessage(uint8_t *buffer,uint16_t *offset,uint8_t ctrlMessage, uint32_t lengthLeft)
    {
        buffer[(*offset)] = ctrlMessage; *offset += 1;
        if(lengthLeft > 128)
        {
            buffer[(*offset)] = (0x80 | (lengthLeft >> 8 & 0xFF));
            buffer[(*offset)+1] = (lengthLeft & 0xFF);
            *offset += 2;
        }
        else
        {
            buffer[(*offset)] = lengthLeft & 0xFF;
            *offset += 1;
        }
    }

    void appendRawBuffer(
        uint8_t *buffer,
        uint16_t *offset,
        uint8_t *src,
        uint16_t len)
    {
        memcpy(&(buffer[(*offset)]), src, len);
        *offset += len;
    }

    ReturnCodes MQTTClient::init(
        const char *id,
        const char *addr,
        uint16_t port,
        uint16_t bufferSize)
    {
        
        this->bufferSize = bufferSize;
        this->clientID = id;
        this->clientAddr = addr;
        this->clientPort = port;
        // 申请缓冲区作为MQTT报文的存储区
        this->buffer = malloc(bufferSize);
        if(this->buffer == NULL) return RET_MALLOC_ERROR;
        return RET_OK;
    }

    void MQTTClient::deinit()
    {
        free(this->buffer);
        this->buffer = NULL;
    }

    ReturnCodes MQTTClient::startConnection(
        const char *addr,
        uint16_t port,
        uint8_t flags,
        uint16_t keepalive,
        const char *willTopic,
        const char *willMsg,
        const char *username,
        const char *password,
        uint32_t timeout
        )
    {
        int iret;

        // 创建本地socket
        this->socketFd = socket(AF_INET,SOCK_STREAM, 0);
        if(this->socketFd < 0) return RET_SOCKET_ERROR;
        const int reuse_flag = 1;
        setsockopt(this->socketFd,
                    SOL_SOCKET, 
                    SO_REUSEADDR,
                    &reuse_flag,
                    sizeof(reuse_flag)
                    );

        static struct sockaddr_in localAddr;
        localAddr.sin_family = AF_INET;
        localAddr.sin_port = htons(this->clientPort);
        localAddr.sin_addr.s_addr = inet_addr(this->clientAddr);
        
        iret = bind(this->socketFd,
            (struct sockaddr *)&localAddr,
            sizeof(localAddr)
            );
        if(iret != 0) return RET_SOCKET_ERROR;

        // 设置远端地址
        static struct sockaddr_in remoteAddr;
        remoteAddr.sin_family = AF_INET;
        remoteAddr.sin_port = htons(port);
        remoteAddr.sin_addr.s_addr = inet_addr(addr);
        uint8_t *msgBuffer = (uint8_t *)(this->buffer);
        uint16_t msgOffset = 0;
        
        // 计算有效载荷长度
        uint16_t payloadLen = 0;
        payloadLen += strlen(this->clientID) + 2;
        if(flags & CONNECT_WILL_FLAG) {
            payloadLen += strlen(willTopic) + strlen(willMsg) + 4;
        }
        if(flags & CONNECT_USER_NAME_FLAG) payloadLen += strlen(username) +2;
        if(flags & CONNECT_PASSWORD_FLAG) payloadLen += strlen(password) +2;

        // 构建固定报文
        appendFixedMessage(msgBuffer,&msgOffset,CTRL_MESSAGE_CONNECT << 4,payloadLen+10U);

        // 构建可变报文
        appendUTF8String(msgBuffer,&msgOffset,"MQTT");
        appendByte(msgBuffer, &msgOffset, 0x04); // 支持MQTT3.1.1
        appendByte(msgBuffer, &msgOffset, flags);
        appendHalfWord(msgBuffer, &msgOffset, keepalive);
        
        // 构建有效载荷
        appendUTF8String(msgBuffer,&msgOffset,this->clientID);

        if(flags & CONNECT_WILL_FLAG) {
            appendUTF8String(msgBuffer,&msgOffset,willTopic);
            appendUTF8String(msgBuffer,&msgOffset,willMsg);
        }
        if(flags & CONNECT_USER_NAME_FLAG) appendUTF8String(msgBuffer,&msgOffset,username);
        if(flags & CONNECT_PASSWORD_FLAG) appendUTF8String(msgBuffer,&msgOffset,password);

        // 建立连接
        iret = connect(this->socketFd,(struct sockaddr *)&remoteAddr,sizeof(remoteAddr));
        if(iret < 0) return RET_SOCKET_ERROR;
        iret = send(this->socketFd,
            msgBuffer,
            msgOffset,
            0
        );
        if(iret < 0) return RET_SOCKET_ERROR;

        // 等待MQTT服务器发回ACK
        ReturnCodes eret = this->waitConnectAck(timeout);
        if(eret == RET_OK) this->connected = true;
        return eret;
    }

    ReturnCodes MQTTClient::waitConnectAck(uint32_t timeout)
    {
        int iret;
        struct timeval timeoutVal;
        fd_set recvFds;
        uint8_t *msgBuffer = (uint8_t *)this->buffer;
        // 设置超时时间
        timeoutVal.tv_sec = timeout / 1000000;
        timeoutVal.tv_usec = (timeout % 1000000)*1000000;
        // 等待接收报文
        FD_ZERO(&recvFds);
        FD_SET(this->socketFd,&recvFds);
        iret = select(this->socketFd+1,&recvFds,NULL,NULL,&timeoutVal);
        if(iret <= 0) return RET_TIMEOUT;
        else
        {
            iret = recv(this->socketFd,msgBuffer,this->bufferSize,0);
            if(msgBuffer[0] == (CTRL_MESSAGE_CONNACK << 4U) && msgBuffer[1] == 0x02 && iret == 4)
            {
                switch(msgBuffer[3])
                {
                    case 0x00:
                        return RET_OK;
                        break;
                    case 0x01:
                        return RET_SERVER_UNSUPPORTED;
                        break;
                    case 0x02:
                        return RET_SERVER_UNSUPPORTED;
                        break;
                    case 0x03:
                        return RET_SERVER_UNAVAILABLE;
                        break;
                    case 0x04:
                        return RET_SERVER_UNAUTHORIZED;
                        break;
                    case 0x05:
                        return RET_SERVER_UNAUTHORIZED;
                        break;
                    default:
                        return RET_MESSAGE_ERROR;
                }
            }
            else
            {
                return RET_MESSAGE_ERROR;
            }
        }
    }


    void MQTTClient::endConnection()
    {
        const uint8_t endMsg[2] = {CTRL_MESSAGE_DISCONNECT << 4,0x00};
        send(this->socketFd,endMsg,2,0);
        close(this->socketFd);
        this->connected = false;
    }

    ReturnCodes MQTTClient::publish(
        const char *topic,
        uint8_t * payloadData,
        uint16_t payloadLength,
        QoSLevels qosLevel,
        bool retain
        )
    {
        int iret;
        uint8_t *msgBuffer = (uint8_t *)(this->buffer);
        uint16_t msgOffset = 0;
        if(!this->connected) return RET_CONN_NOT_OPENED;
        uint16_t lenOfTopic = strlen(topic);        
        // 构建固定报文
        uint8_t qosBits;
        bool appendMsgID;
        uint16_t msgID = random()%UINT16_MAX;
        switch (qosLevel)
        {
            case QOS_LEVEL_0:
                qosBits = 0 << 1;
                appendMsgID = false;
                break;
            case QOS_LEVEL_1:
                qosBits = 1 << 1;
                appendMsgID = true;
                break;
            case QOS_LEVEL_2:
                qosBits = 2 << 1;
                appendMsgID = true;
                break;
            default:
                break;
        }
        
        appendFixedMessage(
            msgBuffer,
            &msgOffset,
            CTRL_MESSAGE_PUBLISH << 4 | qosBits,
            appendMsgID 
            ?(lenOfTopic + payloadLength + 4)
            :(lenOfTopic + payloadLength + 2)
            );

        // 构建可变报文
        appendUTF8String(msgBuffer, &msgOffset, topic);
        if(appendMsgID) appendHalfWord(msgBuffer, &msgOffset, msgID);
        // 发送有效符合
        appendRawBuffer(msgBuffer,&msgOffset, (uint8_t *)payloadData, payloadLength);
        // 发送数据
        iret = send(this->socketFd,
            msgBuffer,
            msgOffset,
            0
        );
        if(iret < 0) return RET_SOCKET_ERROR;
        if(qosLevel == QOS_LEVEL_0) return RET_OK;
        else if(qosLevel == QOS_LEVEL_1)
        {
            return RET_UNIMPLEMENTED;
        }
        else
        {
            return RET_UNIMPLEMENTED;
        }
    }
} // namespace SimpleMQTT
